package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;


import nlp.langmodel.LanguageModel;

public class SingleTrigramModel implements LanguageModel {

	private static String START = "<S>";
	private static String STOP = "</S>";
	private static double INCREMENT_FACTOR = 1.0;
	private GoodTuringCounter<Trigram<String, String, String>> trigramCounter = new GoodTuringCounter<Trigram<String,String,String>>();
	private Collection<List<String>> validationCollection;
	Trigram<Double, Double, Double> lambdas = null;
	private double trigramTotal;
	
	public SingleTrigramModel(Collection<List<String>> trainingCollection, Collection<List<String>> validationCollection) {
		this.validationCollection = validationCollection;
		
		
		// trigram generation
		for (List<String> sentence: trainingCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			stoppedSentence.add(0, START);
			stoppedSentence.add(0, START);
			stoppedSentence.add(STOP);
			
			Iterator<String> iter = stoppedSentence.iterator();
			String first = iter.next();
			String second = iter.next();
			String third = iter.next();
			
			while (iter.hasNext()) {
				Trigram<String, String, String> trigram = new Trigram<String, String, String>(first, second, third);
				trigramCounter.incrementCount(trigram, INCREMENT_FACTOR);
				first = second;
				second = third;
				if (iter.hasNext()) {
					third = iter.next();
				}
			}
		}
		
		// nomarlizing counter
		trigramCounter.normalize();
		trigramTotal = trigramCounter.totalCount();
		System.out.println("trigramCount = " + trigramTotal);
		
		// train validation set
//		lambdas = resolveLambda();
//		lambdaPair = new Pair<Double, Double>(0.0, 1.0);
//		System.out.println("lambdaPair is " + lambdas);
		System.out.println("LanguageModelTester.calculatePerplexity() is " + LanguageModelTester.calculatePerplexity(this, validationCollection));
	}
	
	private Trigram<Double, Double, Double> resolveLambda() {
		
		return null;
	}
	
	private double getWordProbability(List<String> sentence, int index) {
		if (index <= 1) {
			return 0.0;
		}
		String third = sentence.get(index);
		String second = sentence.get(index - 1);
		String first = sentence.get(index - 2);
		Trigram<String, String, String> key = new Trigram<String, String, String>(first, second, third);
		
		double count;
		if (!trigramCounter.keySet().contains(key)) {
			count = trigramCounter.getZeroCount(); // / trigramCounter.totalCount();
		} else {
			count = trigramCounter.getCount(key);
		}
		return Math.log(count / trigramTotal) / Math.log(2.0);
	}
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		double logProbability = 0.0;
		
		for (int i = 0; i < stoppedSentence.size(); i++) {
			logProbability += getWordProbability(stoppedSentence, i);
		}
//		System.out.println("sentence p is " + logProbability);
		return Math.pow(2.0, logProbability);
	}

	@Override
	public List<String> generateSentence() {
		// TODO Auto-generated method stub
		return null;
	}
	

}
